import cmd
import sys
import os
import functools
import ollama
import json
import fitz  # PyMuPDF
import subprocess
import shlex


# 使用base code中的common 和 pocketflow
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(parent_dir)
from base_code import *
from base_code.common import *

# ==============================================================================
# simple_mas_shell 代码介绍
# 在这个案例中，我们引入两个Agent分别处理命令执行和界面显示（交互逻辑仍由cmd库处理）
#  - BDIAgent： 构建命令的工作流
#  - HelpAgent： 显示帮助窗口
#  - CommandRegistry： 存储当前支持的命令信息
# 我们会继续拓展这个案例，来看如何简化这个交互，使用户能充分利用所有已实现工具
# ==============================================================================

# ==============================================================================
# 1. CommandRegistry 是ToolSystem的继承类，添加对支持的命令列表的显示
# ==============================================================================
class CommandRegistry(ToolSystem):
    def __init__(self):
        super().__init__()
        self.lines = [
            "     欢迎使用 SimpleMAS ReAct v1.0 (ReAct 模式)",
            " ",
            "支持自然语言交互，Agent 将通过 思考->行动->观察 循环解决复杂问题",
            " ",
            "您可以直接下达复杂指令, 例如: ",
            "  '请列出当前目录下的所有pdf文件，然后读取 paper.pdf 的第一页并总结其内容'",
            " ",
            "内置命令:",
            "  list_tools             - 列出所有已注册的工具",
            "  help                   - 显示此帮助信息",
            "  exit / quit            - 退出程序"
        ]

    def generate_tools_description(self) -> str:
        """
        【新功能】动态地从已注册的工具生成用于提示词的描述文本。
        """
        if not self._tools:
            return "没有可用的工具。"

        descriptions = []
        for tool_name, tool in self._tools.items():
            # 为LLM创建更清晰的工具签名
            # 例如: read_pdf(path: str, max_pages: int = 1) - 读取PDF前N页的文本内容
            params = []
            # 假设 `memory` 参数对LLM是透明的，不需要展示
            for arg in tool.required_args:
                if arg != 'memory':
                    params.append(f"{arg}: str")  # 简化处理，假设都是字符串

            if tool.optional_args is not None:
                for arg, default in tool.optional_args.items():
                    params.append(f"{arg}: {type(default).__name__} = {default}")

            param_str = ", ".join(params)
            descriptions.append(f"- {tool.name}({param_str}): {tool.description}")

        return "\n".join(descriptions)


##################################################################
# 全局单例对象：工具系统（环境层）
##################################################################
GLOBAL_TOOL_SYSTEM = CommandRegistry()


##################################################################
# 装饰器：简化工具注册
##################################################################
def tool(name: str, description: str, required_args: List[str], optional_args: Optional[Dict[str, Any]] = None):
    """
    用于注册外部工具函数的装饰器。
    - name: 工具名
    - description: 工具描述
    - required_args: 必须参数列表
    - optional_args: 可选参数字典
    """
    def decorator(func: Callable):
        # 将原函数包装为 Tool 并注册
        new_tool = Tool(
            name=name,
            description=description,
            function=func,
            required_args=required_args,
            optional_args=optional_args
        )
        GLOBAL_TOOL_SYSTEM.register_tool(new_tool)

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)

        return wrapper
    return decorator


# ==============================================================================
# 2. HelpAgent 支持UI绘制
# ==============================================================================
class HelpAgent:
    def __init__(self, registry: CommandRegistry):
        self.command_registry = registry

    def show_intro(self):
        try:
            return create_aligned_intro(self.command_registry.lines)
        except Exception:
            return "\n".join(self.command_registry.lines)

    def show_help(self, command=None):
        if command and self.command_registry.has_tool(command):
            t = self.command_registry.get_tool(command)
            usage = f"{t.name} {' '.join(t.required_args)}"
            return f"命令: {t.name}\n功能: {t.description}\n用法: {usage}"
        else:
            return self.show_intro()


# ==============================================================================
# 3. ReActAgent 的核心组件
# ==============================================================================
class AgentMemory:
    """
    Memory class that stores short and long memory of agents

    """
    def __init__(self):
        self.beliefs = {}
        self.raw_input = None
        self.command = None
        self.raw_arg = None
        self.desire = None       # (tool_name, kwargs_dict)
        self.error_message = None
        self.result = None
        self.registry = None

    def reset_workflow_data(self):
        """在每次运行前重置临时数据。"""
        self.command, self.raw_arg, self.desire, self.error_message, self.result = (None,) * 5


# ==============================================================================
#  ReActAgent
# ==============================================================================
class ReActAgent:
    """
    一个实现了 ReAct (Reason-Act-Observe) 循环的 Agent。
    它不再依赖线性的 pocketflow 工作流，而是通过一个内部循环来解决问题。
    """
    def __init__(self, tool_system: CommandRegistry,
                 model_name: str = "qwen3:8b", #"gemma3:4b",
                 max_iterations: int = 10):
        self.tool_system = GLOBAL_TOOL_SYSTEM
        self.model_name = model_name
        self.max_iterations = max_iterations
        self.ollama_client = ollama.AsyncClient()
        self.memory = AgentMemory()

    def _create_react_prompt(self, user_query: str, history: List[str]) -> str:
        """
        【灵魂所在】创建指导LLM进行ReAct思考的系统提示词。
        """
        # 使用新方法动态生成工具列表
        tools_description = self.tool_system.generate_tools_description()

        # 将历史记录格式化为 "scratchpad"
        scratchpad = "\n".join(history) if history else "无"

        # 这是全新的、为ReAct设计的提示词
        return f"""
        你是一个能够解决复杂问题的智能助手。你的任务是接收用户的请求，并通过一个“思考->行动->观察”的循环来完成它。
        注意用户每次会把上一次已经完成的工作以历史记录的形式提供，请根据已经完成的工作，参照工作流程进行对应的反馈。

        **可用工具:**
        你只能使用以下工具。不要编造任何不存在的工具。
        {tools_description}
        系统命令 (dir, cd, mkdir, ...)

         **重要规则:**
        1.  **语言一致性**: 你的"思考"和"最终答案"必须使用与用户请求相同的语言。如果用户用中文提问，你就必须用中文回答。
        
        **工作流程:**
        1.  **思考 (Thought)**: 你要么调用工具 (tool_call)，要么给出最终答案 (final_answer)
        2.  如果observation结果已经能回答用户问题，请直接输出 final_answer。
        2.  如果需要行动: 选择一个工具并提供必要的参数来执行你的计划。你的输出必须是一个严格的JSON对象，格式如下：
            ```json
            {{
                "thought": "这里是你的思考过程...",
                "tool_call": {{
                    "name": "工具名称",
                    "arguments": {{
                        "参数1": "值1",
                        "参数2": "值2"
                    }}
                }}
            }}
            ```
        4. 输出最终答案时你的输出必须是另一个严格的JSON对象，格式如下：
            ```json
            {{
                "thought": "我已经完成了所有步骤，现在可以给出最终答案了。",
                "final_answer": "这里是给用户的最终回答。"
            }}
            ```
        
        当前任务
        用户请求: "{user_query}“ 我已经完成 {scratchpad}
        请开始你的工作流程。            
        """

    async def run_async(self, user_query: str) -> str:
        """
        执行 ReAct 循环来处理用户查询。
        """
        history = []
        current_query = user_query

        for i in range(self.max_iterations):
            print(f"\n🌀 === [循环 {i + 1}/{self.max_iterations}] ===")

            system_prompt = self._create_react_prompt(user_query, history)

            print("🧠 正在思考下一步...")
            content = ""
            try:
                response = await self.ollama_client.chat(
                    model=self.model_name,
                    messages=[{"role": "system", "content": system_prompt}],
                    format="json"
                )
                content = response["message"]["content"]
                llm_output = json.loads(content)

                thought = llm_output.get("thought", "(无思考)")
                #print(f"提示词：{system_prompt}")
                print(f"🤔 思考: {thought}")
                #print(content)

            except Exception as e:
                error_msg = f"❌ LLM响应错误或JSON解析失败: {e}\n原始响应: {content}"
                print(error_msg)
                history.append(f"Observation: {error_msg}")
                continue  # 继续下一次循环，让Agent尝试自我修复

            # 检查是否是最终答案
            if "final_answer" in llm_output:
                final_answer = llm_output["final_answer"]
                print(f"✅ 任务完成。最终答案: {final_answer}")
                return final_answer

            # 检查是否有工具调用
            if "tool_call" in llm_output:
                tool_call = llm_output["tool_call"]
                tool_name = tool_call.get("name")
                tool_args = tool_call.get("arguments", {})

                if not tool_name or not self.tool_system.has_tool(tool_name):
                    observation = f"错误: 尝试调用一个不存在的工具 '{tool_name}'。"
                else:
                    print(f"🛠️ 行动: 调用工具 `{tool_name}`，参数: {tool_args}")
                    try:
                        tool_func = self.tool_system.get_tool(tool_name)
                        # 注入 memory (虽然现在是全局的，但保持接口一致性)
                        tool_args["memory"] = self.memory

                        maybe_coro = tool_func(**tool_args)
                        if asyncio.iscoroutine(maybe_coro):
                            observation = await maybe_coro
                        else:
                            observation = maybe_coro

                        # 限制观察结果的长度，避免prompt过长
                        if isinstance(observation, str) and len(observation) > 1500:
                            observation = observation[:1500] + "\n... (内容已截断)"

                    except Exception as e:
                        observation = f"❌ 工具 '{tool_name}' 执行失败: {e}"

                print(f"👀 观察: {observation}")

                #print(f"新的提示词:{self._create_react_prompt(user_query, history)}")
                history.append(f"Thought: {thought}")
                history.append(f"Action: 调用 {tool_name}({tool_args})")
                history.append(f"Observation: {observation}")
            else:
                # 如果LLM既没有给最终答案，也没有调用工具
                observation = "错误: LLM的响应格式不正确，既没有 final_answer 也没有 tool_call。"
                print(f"👀 观察: {observation}")
                history.append(f"Observation: {observation}")

        return "🔁 Agent 已达到最大循环次数，但仍未得出最终答案。请尝试更具体地提问或分解任务。"



# ==============================================================================
# 注册工具
# ==============================================================================
@tool("set",  "设置键值", ["memory", "raw_arg"])
async def set_value(memory: AgentMemory, raw_arg: str):
    parts = raw_arg.split(None, 1)
    if len(parts) < 2:
        return "❌ 错误: 'set' 命令需要一个键和一个值。\n   用法: set <key> <value>"
    key, value = parts
    memory.beliefs[key] = value
    return f"✅ '{key}' 已设置为 '{value}'"

@tool("get",  "获取键值", ["memory", "raw_arg"])
async def get_value(memory: AgentMemory, raw_arg: str):
    key = raw_arg.strip()
    if not key:
        return "❌ 错误: 'get' 需要一个键。 用法: get <key>"
    if key in memory.beliefs:
        return f"➡️ {key} = {memory.beliefs[key]}"
    else:
        return f"🔍 未找到键: '{key}'"


@tool("list",  "列出键值", ["memory"])
async def list_beliefs(memory: AgentMemory):
    if not memory.beliefs:
        return "ℹ️ 数据库为空"
    else:
        return "\n".join([f"{k}: {v}" for k, v in memory.beliefs.items()])


@tool("delete",  "删除键值", ["memory", "key"])
async def delete_key(memory: AgentMemory, key: str):
    key = key.strip()
    if not key:
        return "❌ 错误: 'delete' 需要一个键。 用法: delete <key>"
    if key in memory.beliefs:
        del memory.beliefs[key]
        return f"✅ 成功: 键 '{key}' 已被删除。"
    else:
        return f"🔍 未找到键: '{key}'"


@tool("read_pdf", "读取pdf", ["path"], {"max_pages": 1})
def read_pdf_text(memory: AgentMemory, path: str, max_pages: int = 1) -> str:
    """
    读取 PDF 前若干页的文本内容
    """
    if not path:
        return "❌ 错误: 'read_pdf' 需要pdf文件路径"
    if not os.path.exists(path):
        raise FileNotFoundError(f"PDF file not found: {path}")

    if not path.lower().endswith('.pdf'):
        raise ValueError("文件必须是PDF格式")

    text = []
    try:
        with fitz.open(path) as doc:
            # 基础信息
            info = {
                "path": path,
                "pages": doc.page_count,
                "metadata": doc.metadata,
                "chapters": []
            }

            for page_num in range(min(len(doc), max_pages)):
                page = doc.load_page(page_num)
                text.append(page.get_text())
        return "\n".join(text) if text else "PDF页面中未找到可读文本。"
    except Exception as e:
        return f"执行时发生错误: {e}"



@tool("search_pdf", "搜索pdf", ["path", "keyword"])
def search_pdf(memory: AgentMemory, path: str, keyword: str) -> list:
    """
    在 PDF 中搜索关键字，返回匹配的页码和上下文。
    """
    if not os.path.exists(path):
        raise FileNotFoundError(f"PDF file not found: {path}")

    results = []
    with fitz.open(path) as doc:
        for page_num, page in enumerate(doc, start=1):
            text = page.get_text()
            if keyword.lower() in text.lower():
                snippet = text[text.lower().index(keyword.lower())-30:
                              text.lower().index(keyword.lower())+30]
                results.append({"page": page_num, "context": snippet})

    return results

@tool("run_terminal_command", "执行系统终端命令", ["command", "args"])
def run_terminal_command(memory: AgentMemory, command: str, args: str) -> str:
    # 尝试系统命令
    # 判断当前是否为 Windows
    is_windows = (os.name == "nt") or (sys.platform == "win32")

    # 使用 shlex 分割参数
    arg_list = shlex.split(args, posix=not is_windows)
    try:
        result = subprocess.run(
            [command] + arg_list,
            capture_output=True,
            text=True,
            shell=True  # Windows/Linux 兼容
        )
        return result.stdout.strip() or result.stderr.strip()
    except Exception as e:
        return f"执行系统命令出错: {e}"
# =======================================================================
# SimpleAgentShell（应用层）
# =======================================================================
class SimpleAgentShell(cmd.Cmd):
    """
    交互式Shell，通过 asyncio.run() 调用Agent的异步方法完成用户指令。
    """
    help_agent = HelpAgent(GLOBAL_TOOL_SYSTEM)
    cmd_agent = ReActAgent(GLOBAL_TOOL_SYSTEM)
    intro = help_agent.show_intro()
    prompt = "agent> "
    print(cmd_agent._create_react_prompt(" 测试输入 ", []))

    def __init__(self):
        super().__init__()

    def default(self, line):
        """
        所有未知命令，直接送到 Agent
        """
        if not line.strip():
            return
        if line in ["exit", "quit", "EOF"]:
            return self.do_exit(line)
        try:
            # 调用 ReAct Agent 的主循环
            final_result = asyncio.run(self.cmd_agent.run_async(line))
            print("\n" + "=" * 20 + " 最终结果 " + "=" * 20)
            print(final_result)
            print("=" * 50)
        except Exception as e:
            print(f"\n❌ Agent在执行过程中遇到严重错误: {e}")

    def do_list_tools(self, arg):
        print("可用命令:", self.cmd_agent.tool_system.list_tools())

    # --- UI控制相关方法
    def do_exit(self, arg):
        print("\n👋 感谢使用，再见！")
        return True
    def do_quit(self, arg): return self.do_exit(arg)
    def do_EOF(self, arg): print(); return self.do_exit(arg)
    def do_help(self, arg):
        print(self.help_agent.show_help(arg.strip()))

def main():
    """程序主入口。"""
    try:
        SimpleAgentShell().cmdloop()
    except KeyboardInterrupt:
        print("\n👋 检测到中断，程序退出。")
        sys.exit(0)

if __name__ == '__main__':
    main()